#version 440 core

layout(local_size_x=512) in;

layout(binding=0) coherent buffer nodeNeighborhoodBuffer {
	int nodeNeighborhood[];
};

layout(binding=1) coherent buffer pressureBuffer {
	float pressure[];
};

layout(binding=2) coherent buffer nodePositionBuffer {
	uvec2 nodePositions[];
};

layout(binding=0, rg32f) uniform image2D velocityImage;

uniform uint numFemNodes;

void main(void)
{
	uint gid = gl_GlobalInvocationID.x;
	if (gid < numFemNodes)
	{
		ivec2 nodePosition = ivec2(nodePositions[gid]);
		vec2 vel = imageLoad(velocityImage, nodePosition).xy;
		float p = pressure[gid];
		vec2 gradp = vec2(0,0);

		uint gidTimes9 = gid*9;

		int n2 = nodeNeighborhood[gidTimes9+2];
		int n4 = nodeNeighborhood[gidTimes9+4];
		int n6 = nodeNeighborhood[gidTimes9+6];
		int n8 = nodeNeighborhood[gidTimes9+8];
		
		vec2 weights = vec2(0.0,0.0);

		// if air: grad+=p; if fluid: grad+=p-(neighbor's p); if solid: nop
		gradp.x = float(n8 >= -1) * p - float(n8 >= 0) * pressure[n8];
		weights.x = float(n8 >= -1);
		
		gradp.x += float(n4 >= 0) * pressure[n4] - float(n4 >= -1) * p;
		weights.x += float(n4 >= -1);

		gradp.y = float(n2 >= -1) * p - float(n2 >= 0) * pressure[n2];
		weights.y = float(n2 >= -1);
		
		gradp.y += float(n6 >= 0) * pressure[n6] - float(n6 >= -1) * p;
		weights.y += float(n6 >= -1);
		
		gradp /= max(vec2(1.0, 1.0), weights);

		imageStore(velocityImage, nodePosition, vec4(vel+gradp, 1.0, 0.0));
	}
}